#include "Containers/BranchPlotOptions.h"
#include "Functions/String.h"
#include "Functions/TerminalIO.h"
#include "Globals/Settings.h"
#include "TRegexp.h"
#include <iostream>
#include <sstream>
using namespace BOSS_Afterburner::String;
using namespace BOSS_Afterburner::TerminalIO;
using namespace std;

BranchPlotOptions::BranchPlotOptions(const TString& input)
{
  Import(input);
}

BranchPlotOptions::BranchPlotOptions(const YAML::Node& input)
{
  Import(input);
}

/// Set all the branch plot options based on an input line.
/// The syntax for this input line is descripted in the class description for
/// @remark The `TString` argument `input` is not a reference, because it is modified in this function.
bool BranchPlotOptions::Import(TString input)
{
  /// -# Check `input` argument: needs at least **three** semicolumns (`;`).
  Int_t nsemicol{input.CountChar(';')};
  if(ImportWarning(input, "has to contain at least 3 semicolumns (at least 4 argument blocks)",
                   nsemicol < 3))
    return false;
  if(ImportWarning(input, Form("contains %d branches, which is impossible to plot", nsemicol - 2),
                   nsemicol > 5))
    return false;
  Ssiz_t count{0};
  /// -# Insert spaces in all occurences of `;;`, because `TString::Tokenize` will skip such an
  /// argument blocks.
  while(input.Contains(";;"))
    input.Insert(input.Index(";;") + 1, " ");
  if(input.EndsWith(";")) input += " ";
  /// -# Get and set the tree name.
  TString tok;
  Ssiz_t  start{0};
  if(ImportWarning(input, "does not contain a tree name", !Tokenize(input, tok, start)))
    return false;
  SetTreeName(tok);
  ++count;
  /// -# The next arguments will be branches. **If they have been loaded, set the output file
  /// name.*8
  while(count < nsemicol - 1 && Tokenize(input, tok, start))
  {
    if(!AddBranch(tok)) return false;
    ++count;
  }
  /// -# Get cut selection options.
  if(ImportWarning(input, "does not contain a cut selection", !Tokenize(input, tok, start)))
    return false;
  SetCutSelection(tok);
  /// -# Get draw options.
  if(ImportWarning(input, "does not contain draw options", !Tokenize(input, tok, start)))
    return false;
  SetDrawOptions(tok);
  /// -# Set `fIsOK` status to `true`.
  fIsOK = true;
  return true;
}

bool BranchPlotOptions::Import(const YAML::Node& input)
{
  if(input["tree name"].IsDefined()) SetTreeName(input["tree name"].as<string>().c_str());
  const auto& branches = input["axes"];
  if(!branches || branches.size() == 0) return false;
  for(const auto& branch : branches)
    AddBranch(branch);
  SetCutSelection(input);
  SetDrawOptions(input);
  SetOutputFileName(input);
  SetTitle(input);
  return true;
}

/// Attempt to add a `AxisNamed` object to the `fListOfAxes` based on the `input` string.
/// Format of the `input` string should be according to the description of the `BranchPlotOptions`
bool BranchPlotOptions::AddBranch(const TString& input)
{
  auto ncommas = input.CountChar(',');
  if(ncommas == 0 || ncommas == 3)
  {
    fListOfAxes.push_back({input});
    return true;
  }
  ImportWarning(input, Form("--> branch \"%s\" is invalid", input.Data()));
  return false;
}

bool BranchPlotOptions::AddBranch(const YAML::Node& input)
{
  AxisNamed ax(input);
  fListOfAxes.push_back(ax);
  return true;
}

/// Set cut selection string.
/// Format should be according to the `selection` argument of the [TTree::Draw](https://root.cern.ch/doc/master/classTTree.html#a73450649dc6e54b5b94516c468523e45) method.
void BranchPlotOptions::SetCutSelection(const TString& selection)
{
  fCutSelection = selection;
}

void BranchPlotOptions::SetCutSelection(const YAML::Node& input)
{
  const auto& cutsel = input["cut selection"];
  if(!cutsel) return;
  fCutSelection.Clear();
  for(Int_t i = 0; i < cutsel.size(); ++i)
  {
    fCutSelection += cutsel[i].as<string>().c_str();
    if(i != cutsel.size() - 1) fCutSelection += " && ";
  }
}

/// Set draw options and logscale bits based on some `input` line.
void BranchPlotOptions::SetDrawOptions(TString input)
{
  /// <ol>
  /// <li> Ignore upper case.
  input.ToLower();
  /// <li> Set logscale bits if the `input` string contains a substring `"log"`. All characters that
  /// are `x`, `y`, or `z` after each `"log"` substring are read and used to set the corresponding
  /// log bits.
  auto start = input.Index("log");
  for(size_t i = 0; i < fSetLog.size(); ++i)
    fSetLog[i] = false; // reset bits
  if(start > 0)
  {
    TString tok;
    while(input.Tokenize(tok, start, "log"))
    {
      for(int i = 0; i < tok.Length(); ++i)
      {
        UChar_t index = tok[i] - 'x'; // unsiged to simplify next if statement
        if(index >= 3) break;         // if not x, y, or z
        fSetLog[index] = true;
      }
    }
    TRegexp search{"log[xyz]*"};
    while(input.Contains(search))
    {
      input.Remove(input.Index(search), input(search).Length());
    }
  }
  /// <li> Remove leading and trailing whitespace characters from what remains of `input` and use to set `fDrawOption`.
  input       = input.Strip(TString::kBoth, ' ');
  input       = input.Strip(TString::kBoth, '\t');
  fDrawOption = input;
  /// </ol>
}

void BranchPlotOptions::SetDrawOptions(const YAML::Node& input)
{
  const auto& draw = input["draw options"];
  const auto& log  = input["log scale"];
  if(draw.IsDefined()) fDrawOption = draw.as<string>().c_str();
  if(log.IsDefined())
  {
    TString logstr = log.as<string>().c_str();
    logstr.ToLower();
    SetNoLog();
    if(logstr.Contains("x")) fSetLog[0] = true;
    if(logstr.Contains("y")) fSetLog[1] = true;
    if(logstr.Contains("z")) fSetLog[2] = true;
  }
}

void BranchPlotOptions::SetNoLog()
{
  for(size_t i = 0; i < fSetLog.size(); ++i)
    fSetLog[i] = false;
}

/// Set name of the `TTree` and the output file name.
/// If the input string `treename` contains a `">"` sign, everything that comes behind it determines the output file name.
void BranchPlotOptions::SetTreeName(const TString& treename)
{
  fTreeName = treename;
  if(treename.Contains(">"))
  {
    fOutputFileName = treename;
    fTreeName.Remove(treename.First('>'));
    fOutputFileName.Remove(0, treename.First('>') + 1);
    Trim(fTreeName);
    Trim(fOutputFileName);
  }
}

void BranchPlotOptions::SetOutputFileName(const TString& filename)
{
  fOutputFileName = filename;
}

void BranchPlotOptions::SetOutputFileName(const YAML::Node& input)
{
  const auto& name = input["output file"];
  if(!name) return;
  SetOutputFileName(name.as<string>().c_str());
}

void BranchPlotOptions::SetAxisLabels(TH1* hist) const
{
  if(!hist && hist->IsZombie()) return;
  for(Int_t i = 0; i < fListOfAxes.size(); ++i)
  {
    TAxis* axis = nullptr;
    switch(i)
    {
      case 0: axis = hist->GetXaxis(); break;
      case 1: axis = hist->GetYaxis(); break;
      case 2: axis = hist->GetZaxis(); break;
    }
    if(!axis) continue;
    TString title = fListOfAxes[i].GetTitle();
    TString unit  = fListOfAxes[i].GetUnit();
    if(title.EqualTo("") && unit.EqualTo("")) continue;
    if(!unit.EqualTo("")) title += " (" + unit + ")";
    axis->SetTitle(title);
  }
}

void BranchPlotOptions::SetPlotLabel(TH1* hist) const
{
  if(!HasTitle()) return;
  if(!hist && hist->IsZombie()) return;
  hist->SetTitle(GetTitle());
}

// * ========================== * //
// * ------- INFORMATION ------ * //
// * ========================== * //

bool BranchPlotOptions::HasLogScale() const
{
  bool setlog{false};
  for(auto it : fSetLog)
    setlog |= it;
  return setlog;
}

void BranchPlotOptions::Print() const
{
  if(!IsOK())
  {
    cout << "Empty BranchPlotOptions";
    if(fTreeName.Length()) cout << " (" << fTreeName << ")";
    cout << endl;
  }
  else
  {
    cout << "BranchPlotOptions for \"" << fTreeName << "\"" << endl;
    cout << "  output = \"" << OutputFileName() << "\"" << endl;
    cout << "  varexp = \"" << VarExp() << "\"" << endl;
    cout << "  cutsel = \"" << fCutSelection << "\"" << endl;
    cout << "  option = \"" << fDrawOption << "\"" << endl;
    if(HasLogScale()) { cout << "  setlog = " << LogXYZ() << endl; }
    cout << "  Branches:" << endl;
    for(auto& branch : fListOfAxes)
      cout << branch.BuildString() << endl;
  }
}

bool BranchPlotOptions::IsOK() const
{
  return fTreeName.Length() && fListOfAxes.size();
}

string BranchPlotOptions::LogXYZ() const
{
  string logxyz;
  for(int i = 0; i < 3; ++i)
    if(fSetLog[i]) logxyz += Form("%c", i + 'x');
  return logxyz;
}

// * ================================ * //
// * ------- ARGUMENT BUILDERS ------ * //
// * ================================ * //

/// Build a `varexp` that is required by the [`TTree::Draw`](https://root.cern.ch/doc/master/classTTree.html#a73450649dc6e54b5b94516c468523e45) method from the contents of `f`.
/// @return The returned string is a `nullptr` if building a `varexp` string failed.
/// @todo Should be moved to `ChainLoader` so that one has access to the corresponding `TTree`.
const char* BranchPlotOptions::VarExp() const
{
  /// -# Return `nullptr` if `fTreeName` of `fListOfAxes` is empty.
  if(!fTreeName.Length()) return nullptr;
  if(!fListOfAxes.size()) return nullptr;
  /// -# Build string of branch names.
  TString varexp;
  bool    setbinning{true};
  for(auto& axis : fListOfAxes)
  {
    varexp.Prepend(Form("%s:", axis.GetName()));
    setbinning &= axis.HasBinning();
  }
  varexp.Remove(varexp.Length() - 1);
  /// -# Append output histogram name for ROOT. This name has to be unique to avoid overwriting
  /// histograms in memory, hence the terribly long and messy name. The name is also used in the
  /// output file name.
  varexp += ">>";
  varexp += BuildHistName();
  /// -# Append binning **if all branches have bin definitions** (all have `AxisNamed::IsOK` is
  /// `true`).
  if(setbinning)
  {
    varexp += "(";
    for(auto& branch : fListOfAxes)
    {
      varexp += Form("%d,%g,%g,", branch.NBins(), branch.From(), branch.To());
    }
    varexp.Remove(varexp.Length() - 1);
    varexp += ")";
  }
  return Form("%s", varexp.Data());
}

/// Build the original line from which this `BranchPlotOptions` object was created.
/// See `Import` to get an idea of the syntax of that line.
string BranchPlotOptions::BuildOriginalString() const
{
  /// -# Return a `nullptr` if this `BranchPlotOptions` object does not have branches or a tree name.
  if(!fTreeName.Length()) return nullptr;
  if(!fListOfAxes.size()) return nullptr;
  /// -# The line starts the name of the `TTree`.
  string line{fTreeName.Data()};
  /// -# If the `fOutputFileName` is different than the `fTreeName`, append `" > "fOutputFileName`.
  if(!fTreeName.EqualTo(fOutputFileName)) line += Form(" > %s", fOutputFileName.Data());
  line += "; ";
  /// -# Now add the branches.
  for(auto& axis : fListOfAxes)
  {
    line += axis.GetName();
    if(axis.HasBinning()) line += Form(", %d, %g, %g", axis.NBins(), axis.From(), axis.To());
    line += "; ";
  }
  /// -# Add the cut selection.
  line += Form("%s; ", fCutSelection.Data());
  /// -# Add the draw options and a line for the log bits if available.
  line += fCutSelection;
  if(HasLogScale()) line += Form(" log%s", LogXYZ().c_str());
  return line;
}

/// Build the a unique name for an output histogram.
/// The histogram has to be unique to avoid overwriting objects in the `gDirectory` of ROOT, see [TDirectory::Get](https://root.cern.ch/doc/master/classTDirectory.html#a411338189331c0ec67b2c56822b7773a).
string BranchPlotOptions::BuildHistName() const
{
  /// -# Return a `nullptr` if this `BranchPlotOptions` object does not have branches or a tree name.
  if(!fListOfAxes.size()) return "";
  /// -# The histogram name starts the name of the `TTree`.
  string line{Form("hist%d_%s_", fID.GetNanoSec(), fTreeName.Data())};
  /// -# This is followed by the branches, which also get identifiers for the plot range.
  for(auto& axis : fListOfAxes)
  {
    line += axis.GetName();
    if(axis.HasBinning()) line += Form("-%d-%g-%g", axis.NBins(), axis.From(), axis.To());
    line += "_";
  }
  line.pop_back();
  return line;
}

// * ====================== * //
// * ------- HELPERS ------ * //
// * ====================== * //

/// Helper function that prints a warning message if `notOK` is set to `true`.
const bool BranchPlotOptions::ImportWarning(const TString& input, const TString& message,
                                            const bool notOK) const
{
  if(!notOK) return false;
  PrintWarning(Form("BranchPlotOptions input line:\n  \"%s\"\n%s", input.Data(), message.Data()));
  return true;
}

/// Helper function that gets the next argument block up till a semicolumn (`;`), strips that argument block of whitespaces and writes it to `tok`.
const bool BranchPlotOptions::Tokenize(const TString& input, TString& tok, Ssiz_t& start,
                                       const char* delim) const
{
  if(!input.Tokenize(tok, start, delim)) { return false; }
  tok = tok.Strip(TString::kBoth, ' ');
  tok = tok.Strip(TString::kBoth, '\t');
  return true;
}